
/*
 * @Author: AlanLee
 * @Date: 2024-06-28 19:11:08
 * @LastEditors: Sian Li lisian_magic@163.com
 * @LastEditTime: 2024-07-05 16:34:53
 * @FilePath: /hudar/src/hudar/src/RadarClass/Radar/src/Radar.cpp
 * @Description: 雷达主程序
 */
#include "../include/Radar.h"

void Radar::init() {
    RCLCPP_INFO(this->get_logger(), "YOU ARE %s", (ENEMY ? "BLUE" : "RED"));
    this->myFrames.setDepth(FRAME_DEPTH);
    this->logger->info("Radar Init!");
    this->logger->flush();
    // 声明需要参数
    this->declare_parameter("model_path", "None");
    this->declare_parameter("lidar_topic", "/livox/lidar");
    this->declare_parameter("camera_params", "None");
    this->declare_parameter("mapping_file", "None");
    this->declare_parameter("record_path", "None");
    this->coverDepth = false;
    this->_if_record = RECORD_SIGN;
    this->record_path = this->get_parameter("record_path").as_string();
    if (this->record_path == "None") {
        RCLCPP_INFO(this->get_logger(), "Can't record_path");
        exit(-1);
    }
    this->mapping_file = this->get_parameter("mapping_file").as_string();
    if (this->mapping_file == "None") {
        RCLCPP_INFO(this->get_logger(), "Can't mapping_file");
        exit(-1);
    }
    this->_model_path = this->get_parameter("model_path").as_string();
    if (this->_model_path == "None") {
        RCLCPP_INFO(this->get_logger(), "Can't Find Models");
        exit(-1);
    }

    this->_lidar_topic = this->get_parameter("lidar_topic").as_string();
    if (this->_lidar_topic == "/livox/lidar") {
        RCLCPP_INFO(this->get_logger(), "Using Default Lidar Topic!");
    }

    this->_camera_params = this->get_parameter("camera_params").as_string();
    if (this->_camera_params == "None") {
        RCLCPP_INFO(this->get_logger(), "Can't Find Camera Params %s",
                    _camera_params.c_str());
    }
    Matrix<float, 3, 3> K_0;
    Matrix<float, 1, 5> C_0;
    Matrix<float, 4, 4> E_0;
    if (!read_param(this->K_0_Mat, this->C_0_Mat, this->E_0_Mat,
                    this->_camera_params)) {
        RCLCPP_INFO(this->get_logger(), "Can't Find Camera Params");
        return;
    }
    cv2eigen(this->K_0_Mat, K_0);
    pcl::PointCloud<pcl::PointXYZ>::Ptr pc(new pcl::PointCloud<pcl::PointXYZ>);
    cv2eigen(this->C_0_Mat, C_0);
    cv2eigen(this->E_0_Mat, E_0);
    this->depthQueue = std::make_shared<DepthQueue>(K_0, C_0, E_0);
    this->pub_robo =
        this->create_publisher<rm_msg::msg::Locations>("/radar/robots", 10);
    this->pub_image = this->create_publisher<sensor_msgs::msg::Image>(
        "/radar/result_view", 10);
    this->sub_cover = this->create_subscription<std_msgs::msg::Bool>(
        "/gui/cover_depth", 10,
        std::bind(&Radar::coverDepthCallback, this, std::placeholders::_1));
    this->sub_exit = this->create_subscription<std_msgs::msg::Bool>(
        "/gui/exit", 10,
        std::bind(&Radar::exitCallback, this, std::placeholders::_1));
    this->LidarListenerBegin();
    this->myLocation = std::make_shared<Location>();
    this->mapMapping = std::make_shared<MapMapping>();
    this->myLocation->decodeMapPoints(this->mapping_file);

    this->_camera = std::make_shared<HKCam>();
    this->armorDetector = std::make_shared<ArmorDetector>(
        this->_model_path + "Armor.engine", this->_model_path + "Armor.onnx");
    this->armorDetector->accessModelTest();

    this->carDetector = std::make_shared<CarDetector>(
        this->_model_path + "Car.engine", this->_model_path + "Car.onnx");
    this->carDetector->accessModelTest();

    if (!this->armorDetector->initModel()) {
        this->logger->info("ArmorDetector Init Failed!");
        this->logger->flush();
        this->stop();
        return;
    }
    if (!this->carDetector->initModel()) {
        this->logger->info("CarDetector Init Failed!");
        this->logger->flush();
        this->stop();
        return;
    }

    if (!this->mapMapping->_is_pass()) {
        this->logger->info("Locate pick start ...Process");
        Mat rvec, tvec;
        unique_lock<shared_timed_mutex> ulk(myMutex_cameraThread);
        try {
            if (!this->myLocation->locate_pick(this->_camera, ENEMY, rvec, tvec,
                                               this->K_0_Mat, this->C_0_Mat,
                                               this->E_0_Mat)) {
                ulk.unlock();
                return;
            }
            ulk.unlock();
        } catch (const std::exception &e) {
            this->logger->error(e.what());
            return;
        }
        this->mapMapping->push_T(rvec, tvec);
        this->logger->info("Locate pick Done");
        this->logger->flush();
    }

    radar_alive = true;
    _mainProcessLoopWorking = true;
    _videoRecorderLoopWorking = true;
    this->videoRecorder = std::make_shared<VideoRecorder>();
    this->_recorder_block = !this->videoRecorder->init(
        this->record_path.c_str(), VideoWriter::fourcc('m', 'p', '4', 'v'),
        Size(ImageW, ImageH));
    RCLCPP_INFO(this->get_logger(), "Starting Thread");
    this->videoRecoderLoop =
        std::thread(std::bind(&Radar::VideoRecorderLoop, this));
    this->mainLoop = std::thread(std::bind(&Radar::mainProcessLoop, this));
    RCLCPP_INFO(this->get_logger(), "Init Done!");
}
void Radar::mainProcessLoop() {
    RCLCPP_INFO(this->get_logger(), "MainProcessLoop Start");
    while (_mainProcessLoopWorking) {
        auto start_t = std::chrono::system_clock::now().time_since_epoch();

        cv::Mat frame = this->_camera->getFrame();

        if (this->_if_record) {
            Mat record_frame = frame.clone();
            this->myFrames.push(record_frame);
        }

        vector<DetectBox> sepTargets = this->carDetector->infer(frame);
        vector<bboxAndRect> pred =
            this->armorDetector->infer(frame, sepTargets);

        this->drawBbox(sepTargets, frame);
        this->drawArmorsForDebug(pred, frame);

        if (this->coverDepth) {
            shared_lock<shared_timed_mutex> slk_pd(this->myMutex_publicDepth);
            for (int i = 0; i < ImageH; ++i) {
                uchar *data = frame.ptr<uchar>(i);
                int k = 0;
                for (int j = 0; j < ImageW; ++j) {
                    if (fabs(this->publicDepth[i][j]) > Epsilon) {
                        int b, g, r;
                        hsv_to_bgr((this->publicDepth[i][j] / 100.) * 255, 255,
                                   255, b, g, r);
                        data[k] = b;
                        data[k + 1] = g;
                        data[k + 2] = r;
                    }
                    k += 3;
                }
            }
            slk_pd.unlock();
        }

        if (pred.size() != 0) {
            this->armor_filter(pred);
            shared_lock<shared_timed_mutex> slk_pd(this->myMutex_publicDepth);
            if (this->publicDepth.size() == 0) {
                slk_pd.unlock();
                this->logger->info("No Lidar Msg , Return");
            } else {
                this->detectDepth(pred);
                vector<ArmorBoundingBox> IouArmors =
                    this->mapMapping->_IoU_prediction(pred, sepTargets);

                this->detectDepth(IouArmors);
                slk_pd.unlock();
                this->mapMapping->mergeUpdata(pred, IouArmors, this->K_0_Mat,
                                              this->C_0_Mat);
                judge_message myJudge_message;
                myJudge_message.task = 1;
                myJudge_message.loc = this->mapMapping->getloc();
                // TODO: send to judge here
                if (myJudge_message.loc.size() > 0) {
                    rm_msg::msg::Locations locs;
                    for (int i = 0, N = myJudge_message.loc.size(); i < N;
                         ++i) {
                        rm_msg::msg::Location location_msg;
                        location_msg.id = myJudge_message.loc[i].id;
                        location_msg.x = myJudge_message.loc[i].x;
                        location_msg.y = myJudge_message.loc[i].y;
                        locs.locs.emplace_back(location_msg);
                    }
                    this->pub_robo->publish(locs);
                }
            }
        }
        auto end_t = std::chrono::system_clock::now().time_since_epoch();
        char ch[255];
        sprintf(ch, "FPS %d",
                int(std::chrono::nanoseconds(1000000000).count() /
                    (end_t - start_t).count()));
        std::string fps_str = ch;
        cv::putText(frame, fps_str, {10, 50}, cv::FONT_HERSHEY_SIMPLEX, 2,
                    {0, 255, 0}, 3);
        this->mapMapping->_plot_region_rect(this->show_region, frame,
                                            this->K_0_Mat, this->C_0_Mat);
        // RCLCPP_INFO(this->get_logger(), "MainProcessLoop Publishing");
        sensor_msgs::msg::Image::SharedPtr msg_ =
            cv_bridge::CvImage(std_msgs::msg::Header(), "bgr8", frame)
                .toImageMsg();
        pub_image->publish(*msg_.get());
        cv::imshow("test", frame);
        this->logger->flush();
    }
    this->logger->critical("MainProcessLoop Exit");
}
void Radar::VideoRecorderLoop() {
    RCLCPP_INFO(this->get_logger(), "Video Recorder Start");
    while (_videoRecorderLoopWorking) {
        if (this->_if_record && this->myFrames.size() > 0 &&
            !this->_recorder_block) {
            this->videoRecorder->write(this->myFrames.front());
            this->myFrames.pop();
        } else if (!this->_if_record) {
            sleep(1);
        }
    }
    this->videoRecorder->close();

    this->logger->critical("VideoRecorderLoop Exit");
}

void Radar::coverDepthCallback(const std_msgs::msg::Bool::SharedPtr msg) {
    RCLCPP_INFO(this->get_logger(), "CoverDepth %d", msg->data);
    this->coverDepth = msg->data;
}
void Radar::exitCallback(const std_msgs::msg::Bool::SharedPtr msg) {
    RCLCPP_INFO(this->get_logger(), "Exit %d", msg->data);
    this->_videoRecorderLoopWorking = false;
    this->videoRecoderLoop.join();
    this->_mainProcessLoopWorking = false;
    this->mainLoop.join();
    this->stop();
    rclcpp::shutdown();
}

Radar::~Radar() { this->stop(); }

void Radar::drawBbox(vector<DetectBox> &bboxs, Mat &img) {
    for (DetectBox &it : bboxs) {
        cv::rectangle(img, Rect(it.x1, it.y1, it.x2 - it.x1, it.y2 - it.y1),
                      Scalar(0, 255, 0), 2);
    }
}

void Radar::drawArmorsForDebug(vector<ArmorBoundingBox> &armors, Mat &img) {
    for (auto &it : armors) {
        Rect temp = Rect(it.x0, it.y0, it.w, it.h);
        cv::rectangle(img, temp, Scalar(255, 255, 0), 2);
        stringstream ss;
        ss << it.cls << "[Depth]" << it.depth << "[Conf]" << it.conf;
        cv::putText(img, ss.str(),
                    Point2i(int(it.x0 + it.w / 2), int(it.y0 + it.h / 2)),
                    FONT_HERSHEY_SIMPLEX, 1, Scalar(255, 0, 255));
    }
}
void Radar::detectDepth(vector<bboxAndRect> &pred) {
    if (pred.size() == 0)
        return;
    for (size_t i = 0; i < pred.size(); ++i) {
        if (pred[i].armor.x0 > ImageW || pred[i].armor.y0 > ImageH ||
            pred[i].armor.x0 + pred[i].armor.w > ImageW ||
            pred[i].armor.y0 + pred[i].armor.h > ImageH)
            continue;
        vector<float> tempBox;
        float center[2] = {pred[i].armor.x0 + pred[i].armor.w / 2.f,
                           pred[i].armor.y0 + pred[i].armor.h / 2.f};
        for (int j = int(max<float>(center[1] - pred[i].armor.h / 2.f, 0.));
             j < int(min<float>(center[1] + pred[i].armor.h / 2.f, ImageH));
             ++j) {
            for (int k = int(max<float>(center[0] - pred[i].armor.w / 2.f, 0.));
                 k < int(min<float>(center[0] + pred[i].armor.w / 2.f, ImageW));
                 ++k) {
                if (this->publicDepth[j][k] == 0)
                    continue;
                tempBox.emplace_back(this->publicDepth[j][k]);
            }
        }
        float tempDepth = 0;
        for (const auto &jt : tempBox) {
            tempDepth += jt;
        }
        pred[i].armor.depth =
            tempBox.size() != 0 ? tempDepth / tempBox.size() : 0.;
        this->logger->info("Depth: [CLS] " + to_string(pred[i].armor.cls) +
                           " [Depth] " + to_string(pred[i].armor.depth));
    }
}
void Radar::detectDepth(vector<ArmorBoundingBox> &armors) {
    if (armors.size() == 0)
        return;
    for (size_t i = 0; i < armors.size(); ++i) {
        if (armors[i].x0 > ImageW || armors[i].y0 > ImageH ||
            armors[i].x0 + armors[i].w > ImageW ||
            armors[i].y0 + armors[i].h > ImageH)
            continue;
        int count = 0;
        vector<float> tempBox;
        float center[2] = {armors[i].x0 + armors[i].w / 2.f,
                           armors[i].y0 + armors[i].h / 2.f};
        for (int j = int(max<float>(center[1] - armors[i].h / 2.f, 0.));
             j < int(min<float>(center[1] + armors[i].h / 2.f, ImageH)); ++j) {
            for (int k = int(max<float>(center[0] - armors[i].w / 2.f, 0.));
                 k < int(min<float>(center[0] + armors[i].w / 2.f, ImageW));
                 ++k) {
                if (this->publicDepth[j][k] == 0)
                    continue;
                tempBox.emplace_back(this->publicDepth[j][k]);
                ++count;
            }
        }
        int tempNum = 0;
        for (const auto &jt : tempBox) {
            tempNum += jt;
        }
        armors[i].depth = count != 0 ? (float)tempNum / (float)count : 0.;
        this->logger->info("Depth: [CLS] " + to_string(armors[i].cls) +
                           " [Depth] " + to_string(armors[i].depth));
    }
}

void Radar::armor_filter(vector<bboxAndRect> &pred) {
    vector<bboxAndRect> results;
    for (int i = 0; i < int(this->mapMapping->_ids.size()); ++i) {
        int max_id = 0;
        float max_conf = 0.f;
        for (size_t j = 0; j < pred.size(); ++j) {
            if ((int)pred[j].armor.cls == this->ids[i] &&
                pred[j].armor.conf - max_conf > Epsilon) {
                max_id = j;
                max_conf = pred[j].armor.conf;
            }
        }
        if (fabs(max_conf) > Epsilon)
            results.emplace_back(pred[max_id]);
    }
    pred.swap(results);
}
void Radar::drawArmorsForDebug(vector<bboxAndRect> &armors, Mat &img) {
    for (auto &it : armors) {
        Rect temp = Rect(it.armor.x0, it.armor.y0, it.armor.w, it.armor.h);
        cv::rectangle(img, temp, Scalar(0, 255, 0), 2);
        cv::putText(img, to_string(int(it.armor.cls)),
                    cv::Point2i(temp.x, temp.y), cv::FONT_HERSHEY_SIMPLEX, 1,
                    cv::Scalar(0, 0, 255), 3);
    }
}

void Radar::stop() {
    this->armorDetector->unInit();
    this->carDetector->unInit();
    RCLCPP_INFO(this->get_logger(), "Radar Shutdown!");
}

void Radar::LidarListenerBegin() {

    this->sub_lidar = this->create_subscription<sensor_msgs::msg::PointCloud2>(
        this->_lidar_topic, 10,
        std::bind(&Radar::LidarCallBack, this, std::placeholders::_1));
    RCLCPP_INFO(this->get_logger(), "Lidar Inited!");
}

void Radar::LidarCallBack(const sensor_msgs::msg::PointCloud2::SharedPtr msg) {
    // RCLCPP_INFO(this->get_logger(), "points_size(%d,%d)", msg->height,
    //             msg->width);
    pcl::PointCloud<pcl::PointXYZ>::Ptr pc =
        std::make_shared<pcl::PointCloud<pcl::PointXYZ>>();
    pcl::fromROSMsg(*msg, *pc);
    std::vector<std::vector<float>> tempDepth = this->depthQueue->pushback(*pc);
    unique_lock<shared_timed_mutex> ulk(this->myMutex_publicDepth);
    this->publicDepth.swap(tempDepth);
    ulk.unlock();
}